from fastapi import APIRouter
from fastapi import HTTPException, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from core.agent import build_agent
from typing import Optional
import json



chat_router = APIRouter()



class ChatRequest(BaseModel):
    content: str
    session_id: Optional[str] = None


@chat_router.post("/chat")
async def chat(request: Request, chat_request: ChatRequest) -> StreamingResponse:
    app = request.app
    agent = await build_agent(app)
    
    content = chat_request.content
    sessionid = chat_request.session_id
    config = {"configurable": {"thread_id": sessionid}}

    async def generate_stream():
        try:
            async for chunk in agent.astream(
                {"messages": [{"role": "user", "content": content}]}, 
                config,
                stream_mode="messages"
            ):
                print(f"Raw chunk: {chunk}")

                # 处理 LangGraph 的消息流
                if isinstance(chunk, tuple) and len(chunk) == 2:
                    message, metadata = chunk
                    
                    # 提取消息内容
                    if hasattr(message, 'content'):
                        content_text = message.content
                        
                        # 构建标准JSON响应
                        response_data = {
                            "type": "data",
                            "content": content_text,
                            "metadata": {
                                "model": metadata.get('ls_model_name', 'unknown'),
                                "threadId": metadata.get('thread_id'),
                                "temperature": metadata.get('ls_temperature', 0.0),
                                "step": metadata.get('langgraph_step'),
                                "node": metadata.get('langgraph_node'),
                                "provider": metadata.get('ls_provider')
                            }
                        }
                        
                        chunk_data = json.dumps(response_data, ensure_ascii=False)
                        yield f"data: {chunk_data}\n\n"
                
                elif hasattr(chunk, 'content'):
                    # 直接处理消息对象
                    response_data = {
                        "type": "data",
                        "content": chunk.content,
                        "metadata": {
                            "model": "unknown",
                            "role": getattr(chunk, 'role', 'assistant')
                        }
                    }
                    
                    chunk_data = json.dumps(response_data, ensure_ascii=False)
                    yield f"data: {chunk_data}\n\n"
                
                elif isinstance(chunk, dict):
                    # 如果是字典格式
                    response_data = {
                        "type": "data",
                        "content": chunk.get('content', str(chunk)),
                        "metadata": chunk.get('metadata', {})
                    }
                    
                    chunk_data = json.dumps(response_data, ensure_ascii=False)
                    yield f"data: {chunk_data}\n\n"
                
                else:
                    # 其他情况，尝试转换为字符串
                    response_data = {
                        "type": "data",
                        "content": str(chunk),
                        "metadata": {}
                    }
                    
                    chunk_data = json.dumps(response_data, ensure_ascii=False)
                    yield f"data: {chunk_data}\n\n"
                    
        except Exception as e:
            print(f"Stream error: {e}")
            error_data = {
                "type": "error",
                "content": f"处理请求时发生错误: {str(e)}",
                "metadata": {}
            }
            yield f"data: {json.dumps(error_data, ensure_ascii=False)}\n\n"
        finally:
            # 发送结束标记
            yield "data: [DONE]\n\n"

    return StreamingResponse(
        generate_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",
        },
    )